import os
import sys
osa = os.path.abspath
osd = os.path.dirname
cur_dir = osd(osa(__file__))
par_dir = osd(cur_dir)
sys.path.insert(0, par_dir)

import argparse

parser = argparse.ArgumentParser()
parser.add_argument('-c', '--cuda', type=str, default='4', help='CUDA device(s) to use')
parser.add_argument('-p', '--port', type=int, default=20022, help='opened port')
args, _ = parser.parse_known_args()
os.environ['CUDA_VISIBLE_DEVICES'] = args.cuda

from diffusers import FluxKontextPipeline
from util_flux import resize_with_aspect
from MODEL_CKP import FLUX_KONTEXT
import torch

from flask import Flask, request, send_file, jsonify
from flask_cors import CORS
from io import BytesIO
from PIL import Image

app = Flask(__name__)
CORS(app)

# 初始化pipeline（全局单例）
pipe = None

def get_pipeline():
    global pipe
    if pipe is None:
        pipe = FluxKontextPipeline.from_pretrained(FLUX_KONTEXT, torch_dtype=torch.bfloat16)
        pipe.load_lora_weights("/data/models/FLUX.1-Turbo-Alpha")
        pipe.to("cuda")
    return pipe

@app.route('/generate', methods=['POST'])
def generate():
    """
    接收图片和参数，返回生成的图片
    POST form-data:
        - image: 文件
        - prompt: str
        - guidance_scale: float (可选, 默认2.5)
        - steps: int (可选, 默认8)
    """
    if 'image' not in request.files:
        return jsonify({'error': 'Missing image or prompt'}), 400

    file = request.files['image']
    prompt = request.form.get('prompt', 'wear glasses')
    guidance_scale = float(request.form.get('guidance_scale', 2.5))
    steps = int(request.form.get('steps', 8))

    try:
        img = Image.open(file.stream).convert('RGB')
    except Exception as e:
        return jsonify({'error': f'Invalid image: {e}'}), 400

    img = resize_with_aspect('', img_pil=img)
    
    # 获取单例pipeline
    pipe = get_pipeline()
    
    with torch.no_grad():
        result = pipe(
            image=img,
            prompt=prompt,
            prompt_2=prompt,
            height=img.height,
            width=img.width,
            num_inference_steps=steps,
            guidance_scale=guidance_scale,
        ).images[0]

    img_io = BytesIO()
    result.save(img_io, format='PNG')
    img_io.seek(0)
    return send_file(img_io, mimetype='image/png')

if __name__ == "__main__":
    # 预加载模型（可选）
    get_pipeline()
    app.run(host='0.0.0.0', port=args.port, debug=True,use_reloader=False)